package com.feidee.fd.sml.algorithm.component.ml.regression

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.LinearRegression

/**
  * @Author tangjinyuan, songhaicheng
  * @Date 2019/3/26 15:26
  * @Description 　线性回归
  * @Reviewer
  */
case class LinearRegressionParam(
                                  override val input_pt: String,
                                  override val output_pt: String,
                                  override val hive_table: String,
                                  override val flow_time: String,
                                  override val modelPath: String,
                                  override var labelCol: String,
                                  override val featuresCol: String,
                                  override var predictionCol: String,
                                  override val metrics: Array[String],
                                  // Spark treeAggregate 算子的参数，适度增大可以加快训练速度，>= 2，默认 2
                                  aggregationDepth: Int,
                                  // 弹性网络混合参数。[0.0, 1.0]，默认 0.0（对于alpha = 0，惩罚是 L2 惩罚。 对于 alpha = 1，它是 L1 惩罚。 对于（0,1）中的 α，惩罚是 L1 和 L2 的组合）
                                  elasticNet: Double,
                                  // 是否需要计算截距，默认 true
                                  fitIntercept: Boolean,
                                  // 最大迭代数，>= 0，默认 100
                                  maxIter: Int,
                                  // 正则化系数，>= 0，默认 0.0
                                  regParam: Double,
                                  // 优化算法，支持 [auto, l-bfgs, normal]，默认 auto（算法自动选择）
                                  solver: String,
                                  // 是否对特征进行标准化，默认 true
                                  standardization: Boolean,
                                  // 优化算法迭代求解过程的收敛阀值，>= 0，默认 1E-6
                                  tol: Double,
                                  // 样本权重列
                                  weightCol: String
                                ) extends MLParam {
  def this() = this(null, null, null, null, null, "label", "features", "prediction", new Array[String](0),
    2, 0.0, true, 100, 0.0, "auto", true, 1E-6, null)

  override def verify(): Unit = {
    super.verify()
    require(aggregationDepth >= 2,"param aggregationDepth must be greater than 1")
    require(elasticNet >= 0.0 && elasticNet <= 1.0,"param elasticNet's range is [0.0, 1.0]")
    require(maxIter >= 0, "param maxIter can't be negative")
    require(regParam >= 0, "param regParam can't be negative")
    require(tol >= 0,"param tol can't be negative")
    val solvers = Array("auto", "l-bfgs", "normal")
    require(solvers.contains(solver.toLowerCase), s"param solver only accepts ${solvers.mkString("[", ", ", "]")}," +
      s" but has $solver")

  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("aggregationDepth" -> aggregationDepth)
    map += ("elasticNetParam" -> elasticNet)
    map += ("fitIntercept" -> fitIntercept)
    map += ("maxIter" -> maxIter)
    map += ("regParam" -> regParam)
    map += ("solver" -> solver)
    map += ("standardization" -> standardization)
    map += ("tol" -> tol)
    map += ("weightCol" -> weightCol)
    map
  }

}


class LinearRegressionComponent extends AbstractRegressionComponent[LinearRegressionParam] {

  override def setUp(param: LinearRegressionParam): PipelineStage = {
    val lr = new LinearRegression()
      .setAggregationDepth(param.aggregationDepth)
      .setElasticNetParam(param.elasticNet)
      .setFitIntercept(param.fitIntercept)
      .setMaxIter(param.maxIter)
      .setRegParam(param.regParam)
      .setSolver(param.solver)
      .setStandardization(param.standardization)
      .setTol(param.tol)
      .setFeaturesCol(param.featuresCol)
      .setLabelCol(param.labelCol)
      .setPredictionCol(param.predictionCol)

    if (tool.isNotNull(param.weightCol)) {
      lr.setWeightCol(param.weightCol)
    }

    lr
  }

}

object LinearRegressionComponent {
  def apply(paramStr: String): Unit = {
    new LinearRegressionComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    LinearRegressionComponent(args(0))
  }

}
